import os
from tqdm import tqdm
import SimpleITK as sitk
import cv2
import numpy as np
from sklearn.metrics import f1_score
from PIL import Image
from PIL.ImageChops import difference


def conver(img_dir, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    img_list = [i for i in os.listdir(img_dir) if ".nii.gz" in i]
    with tqdm(img_list, desc="conver") as pbar:
        for name in pbar:
            image = sitk.ReadImage(os.path.join(img_dir, name))
            image = sitk.GetArrayFromImage(image)[0]
            image = np.rot90(image, -2)  # 顺时针旋转180°
            image = (image * 255).astype(np.uint8)
            cv2.imwrite(os.path.join(output_dir, name.split(".")[0] + ".png"), image)


def computer_dice(mask_dir, label_dir):
    dice_list = []
    name_list = os.listdir(mask_dir)
    label_list = os.listdir(label_dir)

    with tqdm(name_list, desc="dice") as pbar:
        for index, name in enumerate(pbar):
            mask = cv2.imread(os.path.join(mask_dir, name), cv2.IMREAD_GRAYSCALE)
            label = cv2.imread(os.path.join(label_dir, label_list[index]), cv2.IMREAD_GRAYSCALE)

            if mask.shape != label.shape:
                mask = cv2.resize(mask, (label.shape[1], label.shape[0]))
            mask = (mask / 255).ravel().astype(np.int)
            label = (label / 255).ravel().astype(np.int)
            dice = f1_score(y_true=label, y_pred=mask)
            dice_list.append(dice)
    print(sum(dice_list) * 1.0 / len(dice_list))


def blend_dic(output_images, output_infer):
    name_list = os.listdir(output_images)
    infer_list = os.listdir(output_infer)
    for index, name in enumerate(infer_list):
        im1 = Image.open(output_images + '/' + name_list[index])
        im2 = Image.open(output_infer + '/' + name)

        diff = difference(im2, im1)
        # newimg = Image.blend(im1, im2, alpha=0.5)
        diff.save("./data/" + name)


if __name__ == "__main__":
    img_dir = "nnUNet/nnUNet_raw_data_base/nnUNet_raw_data/Task197_2Dwall/inferTs"
    output_images = "./infer_data"
    infer_dir = "nnUNet/nnUNet_raw_data_base/nnUNet_raw_data/Task197_2Dwall/imagesTs"
    output_infer = "./images_data"

    # conver(img_dir, output_images)
    conver(infer_dir, output_infer)
    # computer_dice(output_images, output_infer)

    # img_dir = "nnUNet/nnUNet_raw_data_base/nnUNet_raw_data/Task197_2Dwall/imagesTs"
    # output_images = "./test_data"
    # conver(img_dir, output_images)
    # dice_list = []
    # name_list = os.listdir(output_images)
    # infer_list = os.listdir(output_infer)
    # print(infer_list)
    # for index, name in enumerate(name_list):
    #     print(name)
    #     infer = infer_list[index]
    #     print(infer)
